from sklearn import manifold
import matplotlib.pyplot as plt
import numpy as np

from har.dataset import *
from util.utils import mkdirs


def t_sne(x, y, sample_rate):
    sample = int(sample_rate * len(y))

    shuffle = np.arange(sample)
    np.random.shuffle(shuffle)

    X = x[shuffle]
    y = y[shuffle]
    tsne = manifold.TSNE(n_components=2, init='pca', random_state=520)
    X_tsne = tsne.fit_transform(X)

    print(f"Original data dimension: {X.shape[-1]}. \nEmbedded data dimension: {X_tsne.shape[-1]}.")
    print(f"X: [N, F]={X.shape}")

    x_min, x_max = X_tsne.min(0), X_tsne.max(0)
    X_norm = (X_tsne - x_min) / (x_max - x_min)
    print(f"tsen shape: {X_norm.shape}")
    return X_norm, y

def split_with_label(x, y):
    labels = {}
    for i in range(len(y)):
    #     label = labels[i]
        if int(y[i]) in labels:
            t = np.vstack((labels[int(y[i])], x[i, :]))
            labels[int(y[i])] = t
        else:
            t = x[i, :]
            labels[int(y[i])] = t
    return labels

def scatter_with_label(x, y, label_dic, title="", savefile=None):
    dic = split_with_label(x, y)
    plt.figure(figsize=(6, 6))
    for i in dic:
         plt.scatter(dic[i][:, 0], dic[i][:, 1], marker="*", color=plt.cm.Set1(i), s=15, alpha=1, label=label_dic[i])
    plt.legend(loc="best")
    plt.xticks([])
    plt.yticks([])
    plt.suptitle(title)
    if savefile is not None:
        savefig(savefile)
    plt.show()

def sample_mean_line(x, y, label_dict, title="", savefile=None ):
    labels = {}
    for i in range(len(x)):
    #     label = labels[i]
        if int(y[i]) in labels:
            t = np.vstack((labels[int(y[i])], x[i, :]))
            labels[int(y[i])] = t
        else:
            t = x[i, :]
            labels[int(y[i])] = t

    x_mean = {}
    plt.figure(figsize=(8, 6))
    colors_map = ['gray', 'red', 'orange', 'green', 'blue', 'darkorchid', 'chocolate', 'pink']
    for index, i in enumerate(labels):
        x = labels[i]
        x_mean[i] = np.sum(x, axis=0) / len(x)
        plt.subplot(f"{len(labels)+1}1{index+1}")
        plt.plot(np.arange(len(x_mean[i])), x_mean[i], linewidth=1, color=colors_map[i], label=label_dict[i])
        # plt.yticks([])
        if index + 1 < len(labels):
            plt.xticks([])
        plt.legend(loc="upper right")
    plt.suptitle(title)
    if savefile is not None:
        savefig(savefile)
    plt.show()

def savefig(file):
    mkdirs(file)
    plt.savefig(file)


class har_data_analysis:    
    def __init__(self):
        pass
